package database

import (
	"abdis/datastruct/dict"
	"abdis/interface/database"
	"abdis/interface/redis"
	"abdis/lib/timewheel"
	"abdis/lib/utils"
	"abdis/redis/protocol"
	"strings"
	"time"
)

const (
	dataDictSize = 1 << 16
	ttlDictSize  = 1 << 10
)

type DB struct {
	index      int
	data       *dict.ConcurrentDict
	ttlMap     *dict.ConcurrentDict
	versionMap *dict.ConcurrentDict

	addAof func(database.CmdLine)

	// 回调函数
	insertCallback database.KeyEventCallback
	deleteCallback database.KeyEventCallback
}

type ExecFunc func(db *DB, args [][]byte) (redis.Reply, bool)

// PreFunc analyses command line when queued command to `multi`
// returns related write keys and read keys
type PreFunc func(args [][]byte) ([]string, []string)

// UndoFunc returns undo logs for the given command line
// execute from head to tail when undo
type UndoFunc func(db *DB, args [][]byte) []database.CmdLine

func makeDB(index int) *DB {
	return &DB{
		data:       dict.MakeConcurrent(dataDictSize),
		ttlMap:     dict.MakeConcurrent(ttlDictSize),
		versionMap: dict.MakeConcurrent(dataDictSize),
		index:      index,
		addAof:     func(line database.CmdLine) {},
	}
}

func (db *DB) RWLocks(writeKeys []string, readKeys []string) {
	db.data.RWLocks(writeKeys, readKeys)
}

func (db *DB) RWUnLocks(writeKeys []string, readKeys []string) {
	db.data.RWUnLocks(writeKeys, readKeys)
}

// Remove 不带锁，需要自行添加锁
func (db *DB) Remove(key string) {
	raw, deleted := db.data.RemoveWithLock(key)
	db.ttlMap.Remove(key)
	taskKey := getExpireTask(key)
	timewheel.Cancel(taskKey)
	if cb := db.deleteCallback; cb != nil {
		entity := &database.DataEntity{}
		if deleted > 0 {
			entity = raw.(*database.DataEntity)
		}
		cb(db.index, key, entity)
	}
}

// Removes 无需自行添加锁
func (db *DB) Removes(keys ...string) (deleted int) {
	deleted = 0
	for _, key := range keys {
		_, exists := db.data.GetWithLock(key)
		if exists {
			db.Remove(key)
			deleted++
		}
	}
	return deleted
}

func getExpireTask(key string) string {
	return "expire:" + key
}

func (db *DB) Expire(key string, expireTime time.Time) {
	db.ttlMap.Put(key, expireTime)
	taskKey := getExpireTask(key)
	timewheel.At(expireTime, taskKey, func() {
		keys := []string{key}

		db.RWLocks(keys, nil)
		defer db.RWUnLocks(keys, nil)

		rawExpireTime, ok := db.ttlMap.Get(key)
		if !ok {
			return
		}
		expireTime, _ = rawExpireTime.(time.Time)
		if time.Now().After(expireTime) {
			db.Remove(key)
		}
	})
}

func (db *DB) IsExpired(key string) bool {
	rawExpireTime, ok := db.ttlMap.Get(key)
	if !ok {
		return false
	}
	expireTime, _ := rawExpireTime.(time.Time)
	expired := time.Now().After(expireTime)
	if expired {
		db.Remove(key)
	}
	return expired
}

func (db *DB) Persist(key string) {
	db.ttlMap.Remove(key)
	timewheel.Cancel(getExpireTask(key))
}

/* ---- 数据处理 ---- */

func (db *DB) GetEntity(key string) (*database.DataEntity, bool) {
	raw, ok := db.data.GetWithLock(key)
	if !ok {
		return nil, false
	}
	if db.IsExpired(key) {
		return nil, false
	}
	entity, _ := raw.(*database.DataEntity)
	return entity, true
}

// PutEntity 返回的是新增的数据量
func (db *DB) PutEntity(key string, entity *database.DataEntity) int {
	ret := db.data.PutWithLock(key, entity)
	if cb := db.insertCallback; ret > 0 && cb != nil {
		cb(db.index, key, entity)
	}
	return ret
}

func (db *DB) PutIfAbsent(key string, entity *database.DataEntity) int {
	ret := db.data.PutIfAbsentWithLock(key, entity)
	if cb := db.insertCallback; ret > 0 && cb != nil {
		cb(db.index, key, entity)
	}
	return ret
}

func (db *DB) addVersion(keys ...string) {
	for _, key := range keys {
		code := db.GetVersion(key)
		db.versionMap.Put(key, code+1)
	}
}

func (db *DB) GetVersion(key string) uint32 {
	entity, ok := db.versionMap.Get(key)
	if !ok {
		return 0
	}
	return entity.(uint32)
}

/* ---- 执行命令 ---- */

func validateArity(arity int, cmdArgs [][]byte) bool {
	argNum := len(cmdArgs)

	if arity >= 0 {
		return argNum == arity
	}

	return argNum >= -arity
}

// Exec 执行命令
func (db *DB) Exec(c redis.Connection, cmdLine [][]byte) redis.Reply {
	name := strings.ToLower(string(cmdLine[0]))

	switch name {
	case "multi":
		if len(cmdLine) != 1 {
			return protocol.MakeArgNumErrReply(name)
		}
		return StartMulti(c)
	case "discard":
		if len(cmdLine) != 1 {
			return protocol.MakeArgNumErrReply(name)
		}
		return DiscardMulti(c)
	case "exec":
		if len(cmdLine) != 1 {
			return protocol.MakeArgNumErrReply(name)
		}
		return execMulti(db, c)
	case "watch":
		if !validateArity(-2, cmdLine) {
			return protocol.MakeArgNumErrReply(name)
		}
		return Watch(db, c, cmdLine[1:])
	}

	if c != nil && c.InMultiState() {
		return EnqueueCmd(c, cmdLine)
	}

	return db.execNormalCommand(cmdLine)
}

func (db *DB) execNormalCommand(cmdLine [][]byte) redis.Reply {
	name := string(cmdLine[0])
	cmd := getCmd(string(cmdLine[0]))
	if cmd == nil {
		return protocol.MakeErrReply("ERR unknown command: " + string(cmdLine[0]))
	}
	if !validateArity(cmd.arity, cmdLine) {
		return protocol.MakeArgNumErrReply(string(cmdLine[0]))
	}

	write, read := cmd.prepare(cmdLine[1:])

	db.addVersion(write...)
	db.RWLocks(write, read)
	defer db.RWUnLocks(write, read)

	reply, ok := cmd.executor(db, cmdLine[1:])
	if ok && cmd.IsWriter() && !cmd.AboutExpire() {
		// 失败不写
		// 读命令不写
		// Expire 命令 自己实现
		db.addAof(utils.ToCmdLineRaw(name, cmdLine[1:]...))
	}
	return reply
}

// for transaction
func (db *DB) execWithLock(cmdLine [][]byte) redis.Reply {
	name := strings.ToLower(string(cmdLine[0]))
	cmd := getCmd(name)
	if cmd == nil {
		return protocol.MakeErrReply("ERR unknown command")
	}
	if !validateArity(cmd.arity, cmdLine) {
		return protocol.MakeArgNumErrReply(name)
	}
	reply, _ := cmd.executor(db, cmdLine[1:])
	return reply
}

func (db *DB) ForEach(consumer database.DbConsumer) {
	db.data.ForEach(func(key string, val any) bool {
		entity, ok := val.(*database.DataEntity)
		if !ok {
			return true
		}
		var expireAt *time.Time
		rawExpire, ok := db.ttlMap.Get(key)
		if ok {
			t, _ := rawExpire.(time.Time)
			expireAt = &t
		}
		return consumer(key, entity, expireAt)
	})
}
